#!/usr/bin/env python2
#coding=utf-8

import commands
import os
import sys
import socket
import time
import subprocess
import fcntl
import types
import errno
import getopt
import re
import math
import random
import uuid
import json
from optparse import OptionParser

from config import Config
from utils import Exp, _dmsg, _dwarn, _derror, _get_value, _str2dict, _str2list, _exec_pipe, _exec_shell, _exec_big


class StorageTool(object):
    def __init__(self, config):
        assert(isinstance(config, Config))
        self.config = config
        self.hostname = config.hostname
        self.sysroot = 'system'
        self.lich_bd = os.path.join(self.config.lich, "libexec/lichbd")
        self.lich_inspect = os.path.join(self.config.lich, "libexec/lich.inspect")
        self.lich_admin = os.path.join(self.config.lich, "libexec/lich.admin")

    def list_node(self, skip_offline = True):
        res = _exec_pipe([self.lich_admin, '--listnode', '-v'], 3, False, 15)[:-1]
        d = {} if (len(res) == 0) else _str2dict(res)
        return { k:{'role':v} for (k, v) in d.items() if not (skip_offline and v == 'stopped')}

    def _is_vol(self, path):
        res = _exec_pipe([self.lich_bd, 'vol', 'info', path], 3, False, 15)[:-1]
        return True if 'chkid : vol.' in res else False

    def _is_dir(self, path):
        pass

    def _vol_info(self, path):
        res = _exec_pipe([self.lich_bd, 'vol', 'info', path], 3, False, 15)[:-1]
        if 'chkid : vol.' not in res:
            return None
        voldic = _str2dict(res)
        return voldic

    def _chunk_location(self, chkid):
        """
        'chunk vol.1036.0 info_version 5 @ [match1.rack1.node156:clean match1.rack1.node154:clean match1.rack1.node155:clean ]'

        :param chkid:
        :return:
        """
        res = _exec_pipe([self.lich_inspect, '--chunkinfo', chkid], 3, False, 15)[:-1]
        if chkid not in res:
            return []

        l = res.split('[')
        s = l[len(l) - 1]
        replica = s.strip(']').strip(' ').split(' ')
        return [x.split(':')[0] for x in replica if x]

        # replica = "\n".join([line.strip(' ').strip('*') for line in res.splitlines()][2:])
        # return [k for (k, v) in _str2dict(replica).items() if v == chkid]

    def _find(self, path = '/', recursive = True):
        res = _exec_pipe([self.lich_bd, 'find', path, '-r'], 3, False, 25)[:-1]
        return [line for line in res.splitlines()]

    def list_pool(self):
        res = _exec_pipe([self.lich_admin, '--poolist'], 3, False, 15)[:-1]
        return [] if len(res) == 0 else _str2list(res)

    def list_volume(self):
        res = self._find()
        vol = [r for r in res if self.sysroot not in r and len(r.replace('/', ' ').split()) > 1]

        dic = {}
        for v in vol:
            try:
                info = self._vol_info(v)
                if info:
                    info['replica'] = self._chunk_location(info['chkid'])
                    dic[v] = info
            except Exp, e:
                if e.errno == errno.ENOSPC:
                    continue
                else:
                    raise
        return dic

    def move_chunk(self, chkid, nodes):
        print 'vol:' + chkid + ' move ' + ','.join(nodes)
        _exec_pipe([self.lich_inspect, '--chunkmove', chkid, ','.join(nodes)], 3, False, 15)[:-1]

    def cpu_useable(self):
        return self.config.polling_core

    def core_hash(self, chkid):
        return int(chkid.split('.')[1]) % self.cpu_useable()


class VolBlcTool(StorageTool):
    def __init__(self):
        config = Config()
        super(VolBlcTool, self).__init__(config)

    def is_admin(self):
        self.nodelst = self.list_node()
        return True if self.nodelst[self.hostname]['role'] == 'admin' else False

    def nodelst_dump(self, nodelst):
        for k , v in nodelst.items():
            detail = ""
            for i in range(self.cpu_useable()):
                detail += "core%d:%d " %(i, len(v['volumes'][i]))
            print "node:%s volume:%-2d %s" % (k, v['volcount'], detail)

    def scan(self, show=False):
        nodelst = self.list_node()
        volst = self.list_volume()

        for name, args in nodelst.items():
            args['volcount'] = 0
            args['volumes'] = {}
            for core in range(self.cpu_useable()):
                args['volumes'][core] = []

        for (k, v) in volst.items():
            core = self.core_hash(v['chkid'])
            controller = v['location']
            volinfo = {'path':k, 'chknum':v['chknum'], 'volid':v['chkid'],
                    'replica':v['replica'], 'core':core, 'controller':controller}

            nodelst[controller]['volumes'][core].append(volinfo)
            nodelst[controller]['volcount'] += 1

        if show:
            self.nodelst_dump(nodelst)

        return volst, nodelst

    def scatter(self, volcount, nodecount):
        scatter = []
        avg = volcount / nodecount
        for i in range(nodecount):
            scatter.append(avg)

        cursor = 0
        left = volcount % nodecount
        for i in range(left):
            scatter[cursor] += 1
            cursor += 1
        return scatter

    def rack(self, host):
        splt = host.split('.')
        if len(splt) == 3:
            return splt[1]
        elif len(split) == 2:
            return splt[0]
        else:
            return 'default'

    def host2rack(self, hostname):
        splt = hostname.split('.')
        if len(splt) == 3:
            return splt[1]
        elif len(splt) == 2:
            return splt[0]
        else:
            return 'default'

    def need_check_rack(self, nodelst):
        racklst = {}
        for host, info in nodelst.items():
            racklst[self.host2rack(host)] = []

        if len(racklst) < 2:
            return None

        for host, info in nodelst.items():
            racklst[self.host2rack(host)].append(host)

        return racklst

    def get_node(self, racklst, skip_rack):
        for rack, nodes in racklst.items():
            if rack not in skip_rack:
                num = random.randint(0, len(nodes)-1)
                return nodes[num]
        return None

    def replica_move(self, replica, dest, nodelst):
        '''
        重新计算副本分布, 如果副本中已包含目的副本, 直接交换位置
        如果不包含 则替换主副本位置 并将其余副本中与主副本冲突的替换掉
        '''
        rackdict = {}
        for host, info in nodelst.items():
            rackdict[self.host2rack(host)] = []

        for host, info in nodelst.items():
            rackdict[self.host2rack(host)].append(host)

        if dest in replica:
            i = replica.index(dest)
            replica[0], replica[i] = replica[i], replica[0]
        else:
            replica[0] = dest

        racklst = self.need_check_rack(nodelst)
        if racklst:
            for idx1 in range(1, len(replica)):
                rack = self.host2rack(replica[idx1])
                for idx2 in range(len(replica)):
                    if idx2 == idx1:
                        continue
                    if rack == self.host2rack(replica[idx2]):
                        skip = [self.host2rack(replica[i]) for i in range(len(replica))]
                        dst = self.get_node(racklst, skip)
                        if dst:
                            replica[idx1] = dst
                        else:
                            raise Exp(errno.ENOEXEC, "Oops! can not find rack")

    def balance2(self, interval=20, force=False, debug=False):
        if not self.is_admin():
            if not force:
                raise Exp(errno.EPERM, "not admin!")
        '''
        扫描所有节点和卷的列表
        卷 {'path': xx, 'volid':xx, 'replica':[xx, ], 'chknum':xx, 'core':xx}
        节点 {
                'hostname':{
                        'volcount':xx,
                        'role': xx,
                        'volumes': [
                                coreid: volinfo
                            ]
                        }
             }
        '''
        volst, nodelst = self.scan(debug)

        volcount = len(volst)
        nodecount = len(nodelst)
        print "vol:%d node:%d avg:%d polling_core:%d" % (volcount, nodecount,
                volcount/nodecount, self.cpu_useable())

        '''
        根据当前集群已创建的卷的id，
        模拟计算出每个节点上每个core应该对应的卷的数量
        '''
        core = {}
        core_avg = {}
        core_scatter= {}
        for i in range(self.cpu_useable()):
            core[i] = []

        for k, v in volst.items():
            core[self.core_hash(v['chkid'])].append({k:v})

        for id, vol in core.items():
            core_avg[id] = float(len(vol)) / nodecount
            core_scatter[id] = self.scatter(len(vol), nodecount)
            if debug:
                print "core:%d avg:%0.2f scatter:%s" % (id, core_avg[id], str(core_scatter[id]))
        '''
        根据计算出的每个节点上的每个core上应该分布的卷的数量， 随机取出多余的卷
        放入待选列表中
        '''
        alternative = []
        for id in range(self.cpu_useable()):
            cursor = 0
            st = sorted(nodelst.copy().items(), key = lambda d : len(d[1]['volumes'][id]), reverse=True)
            for node in st:
                host = node[0]
                info = node[1]
                scatter = core_scatter[id]
                count = len(info['volumes'][id])
                if debug:
                    print "%s core:%d should:%d count:%d" % (host, id, scatter[cursor], count)
                if count > scatter[cursor]:
                    for i in range(count - scatter[cursor]):
                        num = random.randint(0, len(info['volumes'][id]) - 1)
                        needmove = info['volumes'][id][num]
                        if debug:
                            print "    --select %s" % needmove['volid']
                        alternative.append(needmove)
                        nodelst[host]['volumes'][id].remove(needmove)
                        nodelst[host]['volcount'] -= 1

                cursor += 1

        if debug:
            print "core pick out %s" % (str([info['volid'] for info in alternative]))

        '''
        将待选列表中的卷重新分配位置
        '''
        move = []
        for vol in alternative:
            core = vol['core']
            st = sorted(nodelst.copy().items(), key = lambda d : d[1]['volcount'])
            for node in st:
                found = False
                host = node[0]
                info = node[1]
                if vol['controller'] != host and (len(info['volumes'][core]) < core_avg[core]):
                    found = True
                    vol['dest'] = host
                    move.append(vol)
                    nodelst[host]['volcount'] += 1
                    nodelst[host]['volumes'][core].append(vol)
                    print '%s from %s move to %s (core)' %(vol['volid'], vol['controller'], host)
                    break
            if not found:
                raise Exp(errno.ENOEXEC, "Oops! can not find moveable node (polling_core)")

        '''
        目前所有卷在polling_core上已经平衡， 需要针对节点，
        选择负载较高的节点上的卷移动到负载最低的节点
        '''
        while True:
            st = sorted(nodelst.copy().items(), key = lambda d : d[1]['volcount'])
            if st[-1][1]['volcount'] <= st[0][1]['volcount'] + 1:
                break

            info = st[-1][1]
            found = False
            for core in info['volumes']:
                for vol in info['volumes'][core]:
                    count1 = len(st[0][1]['volumes'][core])
                    count2 = len(st[-1][1]['volumes'][core])
                    if not vol.has_key('dest') and count1 < core_avg[core] and count2 > core_avg[core]:
                        src = st[-1][0]
                        dst = st[0][0]
                        if debug:
                            print "    --node:%s core:%d count:%d avg:%f"% (src, core,
                                    len(st[0][1]['volumes'][core]), core_avg[core])

                        vol['dest'] = dst
                        move.append(vol)

                        nodelst[src]['volumes'][core].remove(vol)
                        nodelst[src]['volcount'] -= 1
                        nodelst[dst]['volumes'][core].append(vol)
                        nodelst[dst]['volcount'] += 1
                        print '%s from %s move to %s (ratio)' %(vol['volid'], vol['controller'], dst)
                        found = True
                        break
                if found:
                    break
            if not found:
                raise Exp(errno.ENOEXEC, "Oops! can not find moveable node (node avg)")
        '''
        将计算出来的需要移动的卷进行迁移
        '''
        for needmove in move:
            self.replica_move(needmove['replica'], needmove['dest'], nodelst)
            self.move_chunk(needmove['volid'], needmove['replica'])
            time.sleep(interval)

        self.nodelst_dump(nodelst)


def usage():
    print 'Usage: lich.balance [options]'
    print 'Options:'
    print '-h, --help  show this help message and exit'
    print '--scan      cluster scan'
    print '--balance   cluster balance'


def main():
    parser = OptionParser()
    parser.add_option("-s", "--scan", action='store_true', default=False, dest="scan", help="scan")
    parser.add_option("-b", "--balance", nargs=1, type='int', default=False, dest="balance", help="balance")
    parser.add_option("-f", "--force", action='store_true', default=False,
            dest="force", help="force")
    parser.add_option("-d", "--debug", action='store_true', default=False,
            dest="debug", help="debug")

    (options, args) = parser.parse_args()
    if not options.scan and not options.balance:
        usage()
        sys.exit(1)

    balancetool = VolBlcTool()

    if options.scan:
        try:
            balancetool.scan(True)
        except Exp, e:
            print e.err
            sys.exit(e.errno)

    if options.balance:
        try:
            balancetool.balance2(options.balance, options.force, options.debug)
        except Exp, e:
            print e.err
            sys.exit(e.errno)


if __name__ == '__main__':
    main()
